#include <iostream>
#include "mt_cnn.h"
#include "opencv2/opencv.hpp"
#include "opencv2/videoio.hpp"

using namespace std;
using namespace cv;

#if 1
int main(int argc, char *argv[])
{
    vector<string> model_file = {
        "../model/det1.prototxt",
        "../model/det2.prototxt",
        "../model/det3.prototxt"
        //            "../model/det4.prototxt"
    };

    vector<string> trained_file = {
        "../model/det1.caffemodel",
        "../model/det2.caffemodel",
        "../model/det3.caffemodel"
        //            "../model/det4.caffemodel"
    };

    MTCNN mtcnn(model_file, trained_file);

    vector<Rect> rectangles;
    string img_path = argv[1];
    Mat img = imread(img_path);

    mtcnn.detection(img, rectangles);
    for (int i = 0; i < rectangles.size(); i++)
    {
        int green = 100; // confidences[i] * 255;
        int red = 200;   //(1 - confidences[i]) * 255;
        rectangle(img, rectangles[i], cv::Scalar(0, green, red), 3);
    }
    cv::imwrite("b.jpg", img);
    std::cout << "Hello, World!" << std::endl;
    return 0;
}
#endif
#if 0
int main()
{
    vector<string> model_file = {
        "../model/det1.prototxt",
        "../model/det2.prototxt",
        "../model/det3.prototxt"
        //            "../model/det4.prototxt"
    };

    vector<string> trained_file = {
        "../model/det1.caffemodel",
        "../model/det2.caffemodel",
        "../model/det3.caffemodel"
        //            "../model/det4.caffemodel"
    };

    MTCNN mtcnn(model_file, trained_file);

    //VideoCapture cap(0);
    VideoCapture cap("/home/wencc/workplace/Downloads/1.mp4");

    //    VideoWriter writer;
    //    writer.open("../result/SuicideSquad.mp4",CV_FOURCC('M', 'J', 'P', 'G'), 25, Size(1280,720), true);

    Mat img;
    int frame_count = 0;
    while (cap.read(img))
    {
        vector<Rect> rectangles;
        vector<float> confidences;
        std::vector<std::vector<cv::Point>> alignment;
        mtcnn.detection(img, rectangles, confidences, alignment);

        for (int i = 0; i < rectangles.size(); i++)
        {
            int green = confidences[i] * 255;
            int red = (1 - confidences[i]) * 255;
            rectangle(img, rectangles[i], cv::Scalar(0, green, red), 3);
            vector<cv::Point> align = alignment[i];
            for (int j = 0; j < align.size(); j++) {
                cv::circle(img, align[j], 2, cv::Scalar(0, 0, 255), -1);
            }
        }

        frame_count++;
        cv::putText(img, std::to_string(frame_count), cvPoint(3, 13),
                    cv::FONT_HERSHEY_COMPLEX_SMALL, 0.8, cvScalar(0, 255, 0), 1, CV_AA);
        //        writer.write(img);
        imshow("Live", img);
        waitKey(1);
    }

    return 0;
}
#endif
